
import glob
import logging
import os
import shutil

import helper


logger = logging.getLogger(__name__)


class Main():

    def __init__(self, args):
        self.args = args

    def start(self):
        class_names = []

        os.makedirs(self.args.out_dir, exist_ok=True)

        if self.args.classes_remap:
            remap_filename, varname = self.args.classes_remap.rsplit(':', maxsplit=1)

            remap_file = helper.dynamic_import('remap_file', remap_filename)
            classes_remap = getattr(remap_file, varname + '_index')
        else:
            raise Exception('must have')

        # print(classes_remap)
        # return

        # remapper = helper.ClassMixin(class_names, classes_remap)

        for label_fullname in glob.glob(os.path.join(self.args.lbl_dir, '*.txt')):
            basename = os.path.basename(label_fullname)

            modified = False
            parser = helper.YoloTxt(label_fullname)
            boxes = []
            for box in parser.boxes:
                class_id = int(box[0])
                if class_id < len(classes_remap):
                    remapped_class_id = classes_remap[class_id]
                    if remapped_class_id != class_id:
                        box = (remapped_class_id, *box[1:])
                        boxes.append(box)
                        modified = True
                        continue
                boxes.append(box)

            out_fullname = os.path.join(self.args.out_dir, basename)

            if not modified:
                shutil.copyfile(label_fullname, out_fullname)
            else:
                with open(out_fullname, 'w') as f:
                    # for class_id, nxywh in boxes:
                    #     f.write('%d %f %f %f %f\n' % (class_id, *nxywh))
                    for box in boxes:
                        f.write(' '.join(map(str, helper.flatten(box))))
                        f.write('\n')


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--lbl-dir', type=str, default=r'D:\static\datasets\DOTA_split\trainsplit\labelTxt-txt')
    parser.add_argument('--out-dir', type=str, default=r'D:\static\datasets\DOTA_split\trainsplit\labelTxt-txt-2.0')
    parser.add_argument('--classes-remap', type=str, default=r'C:\my\ws\play-img-py\labelfmt\remap_data.py:DOTA_split_r3is')
    parser.add_argument('--classes-file', type=str, default=None)
    args = parser.parse_args()
    Main(args).start()
